# utils/data_loader.py
"""
数据加载和预处理工具
"""

import pandas as pd
import numpy as np
from pathlib import Path
from typing import Dict, List, Optional, Tuple
import yfinance as yf
from datetime import datetime, timedelta
import ta

class DataLoader:
    """数据加载器类"""
    
    def __init__(self, data_dir: str = "data"):
        self.data_dir = Path(data_dir)
        self.data_dir.mkdir(exist_ok=True)
    
    def download_stock_data(self, symbol: str, start_date: str, end_date: str, 
                           save: bool = True) -> Optional[pd.DataFrame]:
        """下载股票数据"""
        try:
            print(f"📊 下载 {symbol} 数据...")
            
            data = yf.download(
                symbol,
                start=start_date,
                end=end_date,
                progress=False
            )
            
            if data.empty:
                print(f"❌ {symbol} 数据为空")
                return None
            
            # 重命名列名
            data.columns = [col.lower() for col in data.columns]
            
            # 保存数据
            if save:
                filename = f"{symbol}_{start_date}_{end_date}.csv"
                data.to_csv(self.data_dir / filename)
                print(f"💾 数据已保存: {filename}")
            
            print(f"✅ {symbol} 数据下载完成: {len(data)}条记录")
            return data
            
        except Exception as e:
            print(f"❌ {symbol} 数据下载失败: {e}")
            return None
    
    def load_local_data(self, filename: str) -> Optional[pd.DataFrame]:
        """加载本地数据"""
        file_path = self.data_dir / filename
        
        if not file_path.exists():
            print(f"❌ 文件不存在: {filename}")
            return None
        
        try:
            data = pd.read_csv(file_path, index_col=0, parse_dates=True)
            
            # 确保列名小写
            data.columns = [col.lower() for col in data.columns]
            
            print(f"✅ 本地数据加载成功: {len(data)}条记录")
            return data
            
        except Exception as e:
            print(f"❌ 本地数据加载失败: {e}")
            return None
    
    def add_technical_indicators(self, data: pd.DataFrame) -> pd.DataFrame:
        """添加技术指标"""
        try:
            df = data.copy()
            
            # 移动平均线
            df['sma_5'] = ta.trend.SMAIndicator(df['close'], window=5).sma_indicator()
            df['sma_10'] = ta.trend.SMAIndicator(df['close'], window=10).sma_indicator()
            df['sma_20'] = ta.trend.SMAIndicator(df['close'], window=20).sma_indicator()
            
            # 指数移动平均线
            df['ema_12'] = ta.trend.EMAIndicator(df['close'], window=12).ema_indicator()
            df['ema_26'] = ta.trend.EMAIndicator(df['close'], window=26).ema_indicator()
            
            # MACD
            macd = ta.trend.MACD(df['close'])
            df['macd'] = macd.macd()
            df['macd_signal'] = macd.macd_signal()
            df['macd_histogram'] = macd.macd_diff()
            
            # RSI
            df['rsi'] = ta.momentum.RSIIndicator(df['close'], window=14).rsi()
            
            # 布林带
            bollinger = ta.volatility.BollingerBands(df['close'], window=20)
            df['bb_upper'] = bollinger.bollinger_hband()
            df['bb_middle'] = bollinger.bollinger_mavg()
            df['bb_lower'] = bollinger.bollinger_lband()
            df['bb_width'] = bollinger.bollinger_wband()
            
            # 成交量指标
            df['volume_sma'] = ta.volume.VolumeSMAIndicator(df['close'], df['volume'], window=20).volume_sma()
            
            # 波动率
            df['atr'] = ta.volatility.AverageTrueRange(df['high'], df['low'], df['close'], window=14).average_true_range()
            
            # 价格变化
            df['price_change'] = df['close'].pct_change()
            df['high_low_ratio'] = df['high'] / df['low']
            df['close_open_ratio'] = df['close'] / df['open']
            
            print("✅ 技术指标添加完成")
            return df
            
        except Exception as e:
            print(f"❌ 技术指标计算失败: {e}")
            return data
    
    def split_data(self, data: pd.DataFrame, train_ratio: float = 0.8, 
                   val_ratio: float = 0.1) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:
        """分割数据为训练集、验证集和测试集"""
        total_len = len(data)
        train_len = int(total_len * train_ratio)
        val_len = int(total_len * val_ratio)
        
        train_data = data.iloc[:train_len].copy()
        val_data = data.iloc[train_len:train_len + val_len].copy()
        test_data = data.iloc[train_len + val_len:].copy()
        
        print(f"📊 数据分割:")
        print(f"   训练集: {len(train_data)}条 ({len(train_data)/total_len:.1%})")
        print(f"   验证集: {len(val_data)}条 ({len(val_data)/total_len:.1%})")
        print(f"   测试集: {len(test_data)}条 ({len(test_data)/total_len:.1%})")
        
        return train_data, val_data, test_data
    
    def normalize_data(self, data: pd.DataFrame, method: str = 'minmax') -> pd.DataFrame:
        """数据标准化"""
        try:
            df = data.copy()
            
            # 选择数值列
            numeric_columns = df.select_dtypes(include=[np.number]).columns
            
            if method == 'minmax':
                from sklearn.preprocessing import MinMaxScaler
                scaler = MinMaxScaler()
                df[numeric_columns] = scaler.fit_transform(df[numeric_columns])
                
            elif method == 'standard':
                from sklearn.preprocessing import StandardScaler
                scaler = StandardScaler()
                df[numeric_columns] = scaler.fit_transform(df[numeric_columns])
                
            elif method == 'robust':
                from sklearn.preprocessing import RobustScaler
                scaler = RobustScaler()
                df[numeric_columns] = scaler.fit_transform(df[numeric_columns])
            
            print(f"✅ 数据标准化完成 (方法: {method})")
            return df
            
        except Exception as e:
            print(f"❌ 数据标准化失败: {e}")
            return data
    
    def create_sequences(self, data: pd.DataFrame, lookback: int = 30, 
                        target_col: str = 'close') -> Tuple[np.ndarray, np.ndarray]:
        """创建序列数据"""
        try:
            sequences = []
            targets = []
            
            for i in range(lookback, len(data)):
                # 输入序列
                seq = data.iloc[i-lookback:i].values
                sequences.append(seq)
                
                # 目标值（下一个时间点的收盘价）
                target = data[target_col].iloc[i]
                targets.append(target)
            
            sequences = np.array(sequences)
            targets = np.array(targets)
            
            print(f"✅ 序列数据创建完成:")
            print(f"   序列形状: {sequences.shape}")
            print(f"   目标形状: {targets.shape}")
            
            return sequences, targets
            
        except Exception as e:
            print(f"❌ 序列数据创建失败: {e}")
            return np.array([]), np.array([])

class MultiAssetDataLoader(DataLoader):
    """多资产数据加载器"""
    
    def __init__(self, data_dir: str = "data"):
        super().__init__(data_dir)
        self.asset_data: Dict[str, pd.DataFrame] = {}
    
    def download_multiple_assets(self, symbols: List[str], start_date: str, 
                                end_date: str) -> Dict[str, pd.DataFrame]:
        """下载多个资产数据"""
        print(f"📊 下载多资产数据: {len(symbols)}个标的")
        
        for symbol in symbols:
            data = self.download_stock_data(symbol, start_date, end_date)
            if data is not None:
                self.asset_data[symbol] = data
        
        print(f"✅ 成功下载 {len(self.asset_data)} 个资产数据")
        return self.asset_data
    
    def create_portfolio_data(self, symbols: List[str], weights: Optional[List[float]] = None) -> pd.DataFrame:
        """创建投资组合数据"""
        if weights is None:
            weights = [1.0 / len(symbols)] * len(symbols)
        
        if len(weights) != len(symbols):
            raise ValueError("权重数量必须与资产数量相等")
        
        portfolio_data = {}
        
        # 获取所有资产的收盘价
        for symbol in symbols:
            if symbol in self.asset_data:
                portfolio_data[symbol] = self.asset_data[symbol]['close']
        
        if not portfolio_data:
            print("❌ 没有可用的资产数据")
            return pd.DataFrame()
        
        # 创建投资组合DataFrame
        portfolio_df = pd.DataFrame(portfolio_data)
        
        # 计算加权投资组合价值
        portfolio_df['portfolio_value'] = 0
        for i, symbol in enumerate(symbols):
            if symbol in portfolio_df.columns:
                portfolio_df['portfolio_value'] += portfolio_df[symbol] * weights[i]
        
        # 计算投资组合收益率
        portfolio_df['portfolio_return'] = portfolio_df['portfolio_value'].pct_change()
        
        print(f"✅ 投资组合数据创建完成: {len(portfolio_df)}条记录")
        return portfolio_df
    
    def get_correlation_matrix(self, symbols: List[str]) -> pd.DataFrame:
        """计算资产相关性矩阵"""
        returns_data = {}
        
        for symbol in symbols:
            if symbol in self.asset_data:
                returns_data[symbol] = self.asset_data[symbol]['close'].pct_change().dropna()
        
        if not returns_data:
            print("❌ 没有可用的收益率数据")
            return pd.DataFrame()
        
        returns_df = pd.DataFrame(returns_data)
        correlation_matrix = returns_df.corr()
        
        print("✅ 相关性矩阵计算完成")
        return correlation_matrix

def create_sample_data(symbol: str = "SAMPLE", days: int = 252, 
                      start_price: float = 100.0) -> pd.DataFrame:
    """创建示例数据"""
    print(f"🎲 创建示例数据: {symbol}")
    
    np.random.seed(42)
    dates = pd.date_range('2023-01-01', periods=days, freq='D')
    
    # 生成价格数据
    returns = np.random.normal(0.001, 0.02, days)
    prices = [start_price]
    
    for ret in returns:
        new_price = prices[-1] * (1 + ret)
        prices.append(max(new_price, 1.0))  # 确保价格不为负
    
    # 创建OHLCV数据
    data = []
    for i in range(days):
        close = prices[i + 1]
        open_price = prices[i] * (1 + np.random.normal(0, 0.005))
        high = max(open_price, close) * (1 + abs(np.random.normal(0, 0.01)))
        low = min(open_price, close) * (1 - abs(np.random.normal(0, 0.01)))
        volume = np.random.randint(1000000, 5000000)
        
        data.append({
            'open': open_price,
            'high': high,
            'low': low,
            'close': close,
            'volume': volume
        })
    
    df = pd.DataFrame(data, index=dates)
    print(f"✅ 示例数据创建完成: {len(df)}条记录")
    
    return df

if __name__ == "__main__":
    # 测试数据加载器
    print("🧪 测试数据加载器")
    print("=" * 40)
    
    # 创建数据加载器
    loader = DataLoader()
    
    # 创建示例数据
    sample_data = create_sample_data("TEST", 100)
    
    # 添加技术指标
    enhanced_data = loader.add_technical_indicators(sample_data)
    print(f"📊 增强数据列数: {len(enhanced_data.columns)}")
    
    # 数据分割
    train, val, test = loader.split_data(enhanced_data)
    
    # 创建序列数据
    sequences, targets = loader.create_sequences(enhanced_data, lookback=10)
    
    print("✅ 数据加载器测试完成")
